#include <iostream>
#include <iomanip>
#include <vector>

#include <sds/sds.h>

#include "cluster.h"
#include "segment.h"

#include <cmath>

using namespace std;
using namespace sds;

typedef SDataView<int> SegmentView;
typedef SDataView<Point> PointView;
typedef SDataView<unsigned char> CharView;

SegmentView segment;



//double LIMIT_1 = 0.0001;
//double LIMIT_2 = 0.012;
//double LIMIT_H = 0.20;

int cell = 2;
double LIMIT_1 = 0.00008;
double LIMIT_2 = 0.002;
double LIMIT_H = 0.4;

int VIEW = 2;

int segcnt = 0;

double sqr(double a) { return a * a; }

bool validCell(const PointView& src, int i, int j) {
  for (int x = 0; x < cell; ++x) {
    for (int y = 0; y < cell; ++y) {
      if (((const Point&)src[i * cell + x][j * cell + y]).zero()) {
        return false;
      }
    }  
  }
  return true;
}

Cluster segmentInfo(const PointView& src, int i, int j) {
  Cluster info;
  for (int x = 0; x < cell; ++x) {
    for (int y = 0; y < cell; ++y) {
      info += src[i * cell + x][j * cell + y];
    }  
  }
  return info;
}

void setCell(int i, int j, int s) {
  for (int x = 0; x < cell; ++x) {
    for (int y = 0; y < cell; ++y) {
      segment[i * cell + x][j * cell + y] = s;
    }
  }
} 

void updateInfos(const std::vector<int>& seg, std::vector<Cluster>& info,
		 const std::vector<int>& pseg, std::vector<Cluster>& pinfo,
		 const std::vector<SInteger>& size, 
		 const Cluster& update, int i, int j) {
  for (int d = 0; d <= j; ++d) {
    if (seg[d] == seg[j]) info[d] = update;
  }
  if (i != 0) {
    for (int d = j + 1; (d + 1) * cell <= size[1]; ++d) {
      if (pseg[d] == seg[j]) pinfo[d] = update;
    }
  }
}

bool homogenous(const Cluster& info) {
  for (int k = 0; k < CH; ++k) {
    double mean = info.expected(k);
    double dev = info.deviance(k);
    if (dev / mean >= LIMIT_H) return false;
  }
  return true;
}


bool similar(const Cluster& x, const Cluster& y) {
  const double dist = 0.1;
  Cluster z = x + y;
  for (int k = 0; k < CH; ++k) {
    double ax = (x.covariance(k, k) + dist) * x.size();
    double ay = (y.covariance(k, k) + dist) * y.size();
    double a = ax + ay;

    double bx = ax + x.size() * sqr(x.expected(k) - z.expected(k));
    double by = ay + y.size() * sqr(y.expected(k) - z.expected(k));
    double b = bx + by;


    if (b == 0.0) continue;
    double c1 = sqrt(pow(a / b, (z.size() / 2)));

    if (a == 0.0) continue;
    double c2 = sqrt(pow(ax * z.size() / (a * x.size()), x.size() - 1) * 
		     pow(ay * z.size() / (a * y.size()), y.size() - 1));
    

    if (c1 < LIMIT_1) return false; 
    if (c2 < LIMIT_2) return false;

  }
  return true;
}

double similarity(const Cluster& x, const Cluster& y) {
  return mon_distance_euclid(x.center(), y.center());
}

void Read(const string& filename, PointView& points, vector<SInteger>& size) {
  CharView source[CH];
  SFile file(filename);
  for (int k = 0; k < CH; ++k) {
    source[k] = file.getDataset(k).getDataView<unsigned char>();
  }
  
  size = source[0].getSize();
 
  points = PointView(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      points[i][j] = point;    
    }
  }
}

int main(int argc, const char *argv[]) {
  std::cout << setprecision(3);

  PointView src;
  vector<SInteger> size;
  
  Read(argv[1], src, size);

  segcnt = 0;
  segment = SDataView<int>(size);

  vector<int> pseg(size[1] / cell);
  vector<Cluster> pinfo(size[1] / cell);

  vector<int> seg(size[1] / cell);
  vector<Cluster> info(size[1] / cell);

  for (int i = 0; (i + 1) * cell <= size[0]; ++i) {
    for (int j = 0; (j + 1) * cell <= size[1]; ++j) {
      info[j] = segmentInfo(src, i, j);
      if (!homogenous(info[j]) || !validCell(src, i, j)) {
	seg[j] = 0;
      } else {
	seg[j] = -1;
      }
    }
    for (int j = 0; (j + 1) * cell <= size[1]; ++j) {
      if (seg[j] == 0) {
	setCell(i, j, 0);
	continue;
      }

      bool found = false;
      int bind = 0;
      double value = 0.0;

      if (j > 0 && seg[j - 1] != 0 && similar(info[j], info[j - 1])) {
        found = true;
        bind = -1;
        value = similarity(info[j], info[j - 1]);
      }

      if (i > 0) {
        for (int v = 0; v <= VIEW && (j + v) * cell <= size[1]; ++v) {
          bool fnd = pseg[j + v] != 0;
          double val = 0.0;
          Cluster update = pinfo[j + v];
          for (int w = v; fnd && w >= 0; --w) {
            fnd = similar(update, info[j + w]) && seg[j + w] != 0;
            double dist = similarity(update, info[j + w]);
            if (dist > val) val = dist;
            update += info[j + w];
          }
          if (fnd) {
            if (!found || val < value) {
              value = val;
              bind = v;
              found = true;
            }
          }
        }
      }

      if (!found) {
        seg[j] = ++segcnt;
        setCell(i, j, seg[j]);
      } else {
        if (bind == -1) {
          seg[j] = seg[j - 1];
          Cluster update = info[j] + info[j - 1];
	  updateInfos(seg, info, pseg, pinfo, size, update, i, j);
	  setCell(i, j, seg[j]);
        } else {
          Cluster update = pinfo[j + bind];
          for (int w = bind; w >= 0; --w) {
            update += info[j + w];
            seg[j + w] = pseg[j + bind];
            setCell(i, j + w, pseg[j + bind]);
          }
          updateInfos(seg, info, pseg, pinfo, size, update, i, j);
        }
      }
    }
    for (int x = 0; x < cell; ++x) {
      for (int j = size[1] - size[1] % cell; j < size[1]; ++j) {
	segment[i * cell + x][j] = 0;
      }
    }
    pseg.swap(seg);
    pinfo.swap(info);
  }

  for (int i = size[0] - size[0] % cell; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      segment[i][j] = 0;
    }
  }

  std::cout << segcnt << std::endl;

  SDataView<unsigned char> target(size);

  WriteSegments(segment, target);
  

  SFile trgfile(argv[2], SFile::CREATE);
  trgfile.createDataset("segmented", sds::UINT8, size).
    setDataView<unsigned char>(target);

  return 0;
}
